'''
Created on Dec 1, 2010

@author: mkiyer

This module contains classes and parsers for pipeline configuration parameters
stored in configuration files.  Currently, sample specific parameters are 
stored in python's 'ConfigParser' or 'INI' format, and run-level parameters are 
stored in an XML file. 
'''
import os
import re
import collections
import ConfigParser
import xml.etree.cElementTree as etree

class PipelineConfig(object):
    '''
    this is the base class for pipeline configuration.  different
    pipelines should inherit this class and add pipeline-specific
    parsing
    
    the default pipeline provides the capability to:
    (1) copy files,
    (2) uncompress files
    (3) run the FastQC profiling tool
    '''    
    def from_xml(self, xmlfile):
        tree = etree.parse(xmlfile)        
        root = tree.getroot()
        # directories
        self.output_dir = root.findtext('output_dir')
        self.tmp_dir = root.findtext('tmp_dir') 
        # fastqc tools
        self.fastqc_bin = root.findtext('fastqc_bin')

class ExomePipelineConfig(PipelineConfig):
    '''
    exome-specific pipeline configuration
    '''
    def from_xml(self, xmlfile):
        # parent class init
        super(ExomePipelineConfig, self).from_xml(xmlfile)
        tree = etree.parse(xmlfile)    
        root = tree.getroot()
        # email addresses
        self.email_addresses = []
        for elem in root.findall("email"):
            self.email_addresses.append(elem.text)

        # tool paths
        self.samtools_path = root.findtext('samtools_path')
        self.picard_path = root.findtext('picard_path')
        self.rscript_bin = root.findtext('rscript_bin')

        # bwa params
        elem = root.find('bwa')
        self.bwa_path = elem.get('path')
        self.bwa_use_mem = float(elem.findtext('use_mem'))
        self.bwa_num_cores = int(elem.findtext('num_cores'))
        self.bwa_num_mismatches = int(elem.findtext('num_mismatches'))

        # gatk params
        elem = root.find('gatk')
        self.gatk_path = elem.get('path')
        self.gatk_resource_path = elem.findtext('resource_path')
        self.gatk_use_mem = float(elem.findtext('use_mem'))
        self.gatk_num_cores = int(elem.findtext('num_cores'))
        self.gatk_snps_fdr_threshold = float(elem.findtext('snps_fdr_threshold'))
        self.gatk_recal_anal_dir = elem.findtext('recal_anal_dir')
        self.gatk_temp_dir = elem.findtext('temp_dir')

        # This field need to change and be a field at the sample level
        self.gatk_target_exons = elem.findtext('target_exons')
        
        # species specific configuration 
        # TODO: simple tuple for now.. if becomes more 
        # complicated then upgrade to a class
        Genome = collections.namedtuple('Genome', ['species',
                                                   'ucsc_db',
                                                   'bwa_ref_genome',
                                                   'gatk_ref_genome',
                                                   'snpdb', 
                                                   'indeldb'])
        self.genomes = {}
        for elem in root.findall("species"):
            species = elem.get("name")
            ucsc_db = elem.findtext('ucsc_db')
            bwa_ref_genome = elem.findtext('bwa_ref_genome')
            gatk_ref_genome = elem.findtext('gatk_ref_genome')
            snpdb = elem.findtext('snpdb')
            indeldb = elem.findtext('indeldb')
            self.genomes[species] = Genome(species=species,
                                           ucsc_db=ucsc_db,
                                           bwa_ref_genome=bwa_ref_genome,
                                           gatk_ref_genome=gatk_ref_genome,
                                           snpdb=snpdb,
                                           indeldb=indeldb)


# Define constants used to name output files from various
# stages of the pipeline
FASTQC_DATA_FILE = 'fastqc_data.txt'
FASTQC_REPORT_FILE = 'fastqc_report.html'

class JobConfig(object):
    '''
    Base class to represent information about a sequence data analysis 'job'
    
    The structure includes job, sample, and sequencing metadata sections
    as well as a 'fastq_files' section that points to the location of the
    raw sequences to be analyzed
    
    Specific applications should inherit the base class and add
    application-specific attributes
    '''
    def from_cfg(self, cfg_file, root_dir=None):
        '''
        parses a ConfigParser file and sets up parameter
        attributes
        
        cfg_file: file to parse
        root_dir: (optional) path to root_dir on the local
        machine running the pipeline.  this is added as a
        prefix to all pipeline output paths
        '''
        config = ConfigParser.SafeConfigParser()
        config.read(cfg_file)
        # parse job parameters
        self.name = config.get("job", "name")
        # base output dir
        if root_dir is not None:
            self.output_dir = os.path.join(root_dir, self.name)
        else:
            self.output_dir = self.name
        # final output destination (not local working dir)
        self.dst_output_dir = config.get("job", "outputdir")
        self.remote = config.getboolean("job", "remote")
        self.remote_ip = None
        if config.has_option("job", "remote_ip"):            
            self.remote_ip = config.get("job", "remote_ip")
        # parse sample parameters
        self.run_id = config.get("sample", "run_id")
        self.lib_id = config.get("sample", "lib_id")
        self.sample_id = config.get("sample", "sample_id")
        self.species = config.get("sample", "species")
        self.sample_description = config.get("sample", "description")        
        # parse sequencing parameters
        self.seq_center = config.get("sequencing", "center")
        self.seq_date = config.get("sequencing", "date")
        self.seq_platform = config.get("sequencing", "platform")
        self.read_length = config.getint("sequencing", "read_length")
        self.fragment_layout = config.get("sequencing", "fragment_layout")
        self.fragment_length_mean = config.get("sequencing", "fragment_length_mean")
        self.quality_scores = config.get("sequencing", "quality_scores")
        # parse fastq files
        fastq_list = config.items("fastq_files")
        src_fastq_files = {}
        dst_fastq_files = {}
        fastq_files = {}
        fastqc_data = {}
        fastqc_html = {}
        for mate_string,fastq_file in fastq_list:
            m = re.match(r'mate(\d)', mate_string)
            mate_num = int(m.group(1))
            src_fastq_files[mate_num] = fastq_file
            ext = os.path.splitext(fastq_file)[1]
            dst_fastq_files[mate_num] = os.path.join(self.output_dir, 'mate%d%s' % (mate_num,ext))
            fastq_files[mate_num] = os.path.join(self.output_dir, "mate%d.fq" % (mate_num))            
            fastqc_subdir = os.path.join(self.output_dir, '%s_fastqc' % os.path.basename(fastq_files[mate_num]))
            fastqc_data[mate_num] = os.path.join(fastqc_subdir, FASTQC_DATA_FILE)
            fastqc_html[mate_num] = os.path.join(fastqc_subdir, FASTQC_REPORT_FILE) 
        self.src_fastq_files = [src_fastq_files[mate] for mate in xrange(len(src_fastq_files))]
        self.dst_fastq_files = [dst_fastq_files[mate] for mate in xrange(len(dst_fastq_files))]
        self.fastq_files = [fastq_files[mate] for mate in xrange(len(fastq_files))]
        self.fastqc_data_files = [fastqc_data[mate] for mate in xrange(len(fastqc_data))]
        self.fastqc_html_files = [fastqc_html[mate] for mate in xrange(len(fastqc_data))]


class ExomeJobConfig(JobConfig):
    def from_cfg(self, cfg_file, root_dir=None):
        # parent class init
        super(ExomeJobConfig, self).from_cfg(cfg_file, root_dir)        
        # bwa alignment dir
        self.align_dir = os.path.join(self.output_dir, "align")
        # bwa output SAM file
        self.align_sam_file = os.path.join(self.align_dir, "aligned_reads.sam")
        self.align_bam_file = os.path.join(self.align_dir, "aligned_reads.bam")
        # gatk output dir
        self.gatk_dir = os.path.join(self.output_dir, "gatk")
        self.gatk_recal_dir = os.path.join(self.gatk_dir, "recal_analysis")
        #sample_lanes_dict={'aM18':['/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.bam']}
